"""
SSE (Server-Sent Events) 管理器
实现观察者模式，支持事件广播和客户端管理
"""

import asyncio
from loguru import logger
import uuid
from typing import AsyncGenerator, Any, Optional
from collections import defaultdict
from fastapi import Request

from app.models.sse_message import SSEMessage
from app.models.task.task_data import Task


class SSEClient:
    """SSE客户端"""

    def __init__(self, client_id: str, request: Request):
        self.client_id = client_id
        self.request = request
        self.queue = asyncio.Queue(maxsize=100)  # 消息队列
        self.connected = True
        self.subscribed_events: set[str] = set()

    async def send_message(self, message: SSEMessage) -> bool:
        """发送消息给客户端"""
        if not self.connected:
            return False

        try:
            # 非阻塞发送，避免队列满时阻塞
            self.queue.put_nowait(message)
            return True
        except asyncio.QueueFull:
            logger.warning(f"客户端 {self.client_id} 消息队列已满，断开连接")
            self.connected = False
            return False
        except Exception as e:
            logger.error(f"发送消息到客户端 {self.client_id} 失败: {e}")
            self.connected = False
            return False

    async def disconnect(self):
        """断开连接"""
        self.connected = False
        # 放入结束标记
        try:
            await self.queue.put(None)
        except:
            pass

    def subscribe_event(self, event_type: str):
        """订阅事件类型"""
        self.subscribed_events.add(event_type)

    def unsubscribe_event(self, event_type: str):
        """取消订阅事件类型"""
        self.subscribed_events.discard(event_type)

    def is_subscribed(self, event_type: str) -> bool:
        """检查是否订阅了指定事件"""
        return event_type in self.subscribed_events or len(self.subscribed_events) == 0


class SSEManager:
    """SSE管理器 - 单例模式"""

    _instance = None

    def __new__(cls):
        if cls._instance is None:
            cls._instance = super().__new__(cls)
            cls._clients: dict[str, SSEClient] = {}
            cls._event_subscribers: dict[str, set[str]] = defaultdict(set)
        return cls._instance

    def __init__(self):
        if not hasattr(self, "_initialized"):
            self._clients: dict[str, SSEClient] = {}
            self._event_subscribers: dict[str, set[str]] = defaultdict(set)
            self._initialized = True

    def add_client(self, request: Request) -> str:
        """添加新客户端"""
        client_id = str(uuid.uuid4())
        client = SSEClient(client_id, request)
        self._clients[client_id] = client
        logger.info(f"新客户端连接: {client_id}")
        return client_id

    async def remove_client(self, client_id: str):
        """移除客户端"""
        if client_id in self._clients:
            client = self._clients[client_id]

            # 标记为断开连接
            client.connected = False

            # 清空队列中的未发送消息
            try:
                while not client.queue.empty():
                    try:
                        client.queue.get_nowait()
                    except asyncio.QueueEmpty:
                        break
            except Exception as e:
                logger.debug(f"清理客户端 {client_id} 队列时出错: {e}")

            # 发送断开标记
            try:
                await client.queue.put(None)
            except:
                pass

            # 从事件订阅中移除
            for event_type in list(self._event_subscribers.keys()):
                self._event_subscribers[event_type].discard(client_id)

            # 从客户端列表中移除
            del self._clients[client_id]
            logger.info(f"客户端断开连接: {client_id}")

        # 如果客户端数量为0，触发垃圾回收
        if not self._clients:
            import gc

            gc.collect()

    async def broadcast_message(
        self,
        event_type: str,
        data: dict[str, Any],
        target_clients: Optional[set[str]] = None,
    ):
        """广播消息到所有或指定客户端"""
        if not self._clients:
            return  # 没有客户端时直接返回

        message = SSEMessage(event=event_type, data=data, id=str(uuid.uuid4()))

        clients_to_notify = target_clients or set(self._clients.keys())
        disconnected_clients = []

        # 使用并发发送提升性能
        async def send_to_client(client_id):
            if client_id not in self._clients:
                return False

            client = self._clients[client_id]
            if not client.connected:
                return False

            # 检查客户端是否订阅了此事件
            if not client.is_subscribed(event_type):
                return True

            return await client.send_message(message)

        # 并发发送给所有客户端
        tasks = []
        for client_id in list(clients_to_notify):
            if client_id in self._clients:
                tasks.append((client_id, send_to_client(client_id)))

        if tasks:
            results = await asyncio.gather(
                *[task[1] for task in tasks], return_exceptions=True
            )

            for (client_id, _), result in zip(tasks, results):
                if isinstance(result, Exception) or result is False:
                    disconnected_clients.append(client_id)

        # 清理断开的客户端
        cleanup_tasks = [
            self.remove_client(client_id) for client_id in disconnected_clients
        ]
        if cleanup_tasks:
            await asyncio.gather(*cleanup_tasks, return_exceptions=True)

    async def send_task_update(self, task: Task):
        """发送任务更新消息"""
        await self.broadcast_message(
            event_type="task_update",
            data={
                "task_id": task.id,
                "name": task.name,
                "status": task.status.value,
                "progress": getattr(task, "progress", None),
                "result": task.result,
                "error": task.error,
                "updated_at": (
                    task.completed_at.isoformat() if task.completed_at else None
                ),
            },
        )

    async def send_progress_update(
        self,
        task_id: str,
        task_name: str,
        progress: float,
        message: Optional[str] = None,
    ):
        """发送任务进度更新"""
        await self.broadcast_message(
            event_type="progress_update",
            data={
                "task_id": task_id,
                "task_name": task_name,
                "progress": progress,
                "message": message,
                "timestamp": asyncio.get_event_loop().time(),
            },
        )

    async def send_system_message(self, message: str, level: str = "info"):
        """发送系统消息"""
        await self.broadcast_message(
            event_type="system_message",
            data={
                "message": message,
                "level": level,
                "timestamp": asyncio.get_event_loop().time(),
            },
        )

    async def get_client_stream(self, client_id: str) -> AsyncGenerator[str, None]:
        """获取客户端消息流"""
        if client_id not in self._clients:
            return

        client = self._clients[client_id]

        # 发送连接确认消息
        welcome_message = SSEMessage(
            event="connected", data={"client_id": client_id, "message": "连接成功"}
        )
        await client.send_message(welcome_message)

        try:
            while client.connected:
                try:
                    # 等待消息，设置超时避免长时间阻塞
                    message = await asyncio.wait_for(client.queue.get(), timeout=30.0)

                    if message is None:  # 结束标记
                        break

                    yield message.to_sse_format()

                except asyncio.TimeoutError:
                    # 发送心跳消息
                    heartbeat = SSEMessage(
                        event="heartbeat",
                        data={"timestamp": asyncio.get_event_loop().time()},
                    )
                    yield heartbeat.to_sse_format()

                except Exception as e:
                    logger.error(f"客户端 {client_id} 消息处理错误: {e}")
                    break

        except Exception as e:
            logger.error(f"客户端 {client_id} 流处理错误: {e}")

        finally:
            await self.remove_client(client_id)

    def subscribe_client_to_event(self, client_id: str, event_type: str):
        """订阅客户端到特定事件"""
        if client_id in self._clients:
            self._clients[client_id].subscribe_event(event_type)
            self._event_subscribers[event_type].add(client_id)

    def get_clients_count(self) -> int:
        """获取连接的客户端数量"""
        return len(self._clients)

    async def shutdown(self):
        """关闭SSE管理器，清理所有客户端"""
        logger.info(f"正在关闭SSE管理器，当前有 {len(self._clients)} 个客户端连接")

        # 获取所有客户端ID的副本
        client_ids = list(self._clients.keys())

        # 并发断开所有客户端连接
        disconnect_tasks = []
        for client_id in client_ids:
            disconnect_tasks.append(self.remove_client(client_id))

        if disconnect_tasks:
            try:
                await asyncio.wait_for(
                    asyncio.gather(*disconnect_tasks, return_exceptions=True),
                    timeout=3.0,
                )
            except asyncio.TimeoutError:
                logger.warning("断开客户端连接超时")

        # 清理订阅者
        self._event_subscribers.clear()

        # 确保客户端字典为空
        self._clients.clear()

        logger.info("SSE管理器已关闭")

    def get_clients_info(self) -> dict[str, Any]:
        """获取客户端信息"""
        return {
            "total_clients": len(self._clients),
            "clients": [
                {
                    "client_id": client_id,
                    "connected": client.connected,
                    "subscribed_events": list(client.subscribed_events),
                }
                for client_id, client in self._clients.items()
            ],
        }

    async def get_task_progress_stream(self, request: Request, task_id: str):
        """为特定任务创建SSE流"""
        client_id = f"task_stream_{task_id}_{id(request)}"
        client = SSEClient(client_id, request)
        client.subscribe_event(f"task_{task_id}")
        client.subscribe_event("task_progress")
        self._clients[client_id] = client

        async def event_stream():
            try:
                # 发送连接确认
                welcome = SSEMessage(
                    event="connected",
                    data={
                        "client_id": client_id,
                        "task_id": task_id,
                        "message": f"已连接到任务 {task_id} 的进度流",
                    },
                )
                yield welcome.to_sse_format()

                while client.connected:
                    # 检查客户端是否断开连接
                    if await request.is_disconnected():
                        break

                    try:
                        # 等待消息
                        message = await asyncio.wait_for(
                            client.queue.get(), timeout=30.0
                        )
                        if message is None:  # 结束标记
                            break
                        yield message.to_sse_format()
                    except asyncio.TimeoutError:
                        # 发送心跳
                        heartbeat = SSEMessage(
                            event="heartbeat",
                            data={"timestamp": asyncio.get_event_loop().time()},
                        )
                        yield heartbeat.to_sse_format()

            except Exception as e:
                logger.error(f"Error in task progress stream: {e}")
            finally:
                await self.remove_client(client_id)

        return event_stream()

    async def get_all_progress_stream(self, request: Request):
        """获取所有任务进度的SSE流"""
        client_id = f"all_progress_{id(request)}"
        client = SSEClient(client_id, request)
        client.subscribe_event("task_progress")
        client.subscribe_event("system_message")
        self._clients[client_id] = client

        async def event_stream():
            try:
                # 发送连接确认
                welcome = SSEMessage(
                    event="connected",
                    data={"client_id": client_id, "message": "已连接到所有任务进度流"},
                )
                yield welcome.to_sse_format()

                while client.connected:
                    if await request.is_disconnected():
                        break

                    try:
                        message = await asyncio.wait_for(
                            client.queue.get(), timeout=30.0
                        )
                        if message is None:
                            break
                        yield message.to_sse_format()
                    except asyncio.TimeoutError:
                        heartbeat = SSEMessage(
                            event="heartbeat",
                            data={"timestamp": asyncio.get_event_loop().time()},
                        )
                        yield heartbeat.to_sse_format()

            except Exception as e:
                logger.error(f"Error in all progress stream: {e}")
            finally:
                await self.remove_client(client_id)

        return event_stream()


# 全局实例
sse_manager = SSEManager()
